import matplotlib.pyplot as plt
import numpy as np
import offline_keypoint_detect_utils as utils
import models.barcode_model_small as models
import torch

if __name__ == "__main__":
    postprocess_params = utils.PostProcessParams(
        width=128,
        height=128,
        heatValueThreshold=0.5,
        distThreshold=10.0,
        maxDistanceDifference=30,
        angleDifferenceThreshold=round(np.pi / 4, 2),
    )

    img_path = r"C:\Users\Administrator\Desktop\barcode_train\real_barcode_data\0.png"
    # mdoel_path = r"C:\Users\Administrator\Desktop\aidc\keypoint_detection\train_cxf_18\sigma_2.0epoch_16\model.pt"
    model_path = r"C:\Users\Administrator\Desktop\barcode_train\runs_barcode_2d\runs1\sigma_2epoch_40.pt"
    model = models.GrayDemoNet2()
    device = torch.device("cpu")
    model.to(device)
    checkpoint = torch.load(model_path, map_location=device)
    if isinstance(checkpoint, dict):
        model.load_state_dict(checkpoint["model_state_dict"])
    else:
        model = checkpoint
    img, outputs_np = utils.detect_image(img_path, model, postprocess_params, 128, 128)

    plt.subplot(2, 2, 1)
    plt.title("heatmap")
    plt.imshow(outputs_np[0, 0, :, :], cmap="gray")
    plt.subplot(2, 2, 2)
    plt.title("sin")
    plt.imshow(outputs_np[0, 1, :, :].squeeze(), cmap="gray")
    plt.subplot(2, 2, 3)
    plt.title("cos")
    plt.imshow(outputs_np[0, 2, :, :].squeeze(), cmap="gray")
    plt.subplot(2, 2, 4)
    plt.title("Visualization")
    plt.imshow(img)
    plt.savefig("analyse.png", dpi=300)
    plt.show()
